import subprocess
import platform
import os
import sys
import time

MODELS = [
    ("deepseek-r1:1.5b", "1.1GB", "推荐"),
    ("deepseek-r1:7b", "4.7GB", "推荐"),
    ("deepseek-r1:8b", "4.9GB", "推荐"),
    ("deepseek-r1:14b", "9GB", "推荐"),
    ("deepseek-r1:32b", "20GB", "不推荐"),
    ("deepseek-r1:70b", "43GB", "不推荐"),
    ("deepseek-r1:671b", "404GB", "不推荐"),
]

def check_ollama_installed():
    try:
        result = subprocess.run(
            ["ollama", "--version"],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )
        return result.returncode == 0
    except FileNotFoundError:
        return False

def download_ollama():
    system = platform.system()
    if system == "Windows":
        url = "https://ghproxy.cfd/https:/github.com/ollama/ollama/releases/download/v0.6.8/OllamaSetup.exe"
        path = os.path.join(os.getenv("TEMP"), "OllamaSetup.exe")
    elif system == "Linux":
        url = "https://ollama.com/install.sh"
        path = "/tmp/install.sh"
    elif system == "Darwin":
        url = "https://github.moeyy.xyz/https://github.com/ollama/ollama/releases/download/v0.5.7/Ollama-darwin.zip"
        path = "/tmp/Ollama-darwin"
    else:
        print("不支持的操作系统")
        sys.exit(1)

    print(f"正在下载Ollama安装包...")
    try:
        if system == "Windows":
            subprocess.run(f"curl -Lo {path} {url}", check=True, shell=True)
        else:
            subprocess.run(f"curl -fsSL {url} -o {path}", check=True, shell=True)
    except subprocess.CalledProcessError:
        print("下载失败，请检查网络连接")
        sys.exit(1)
    return path

def install_ollama():
    system = platform.system()
    install_path = download_ollama()
    
    print("正在安装Ollama...")
    try:
        if system == "Windows":
            subprocess.run([install_path, "/S"], check=True)
        elif system == "Linux":
            subprocess.run(f"chmod +x {install_path} && sudo bash {install_path}", shell=True, check=True)
        elif system == "Darwin":
            subprocess.run(f"chmod +x {install_path} && sudo {install_path}", shell=True, check=True)
    except subprocess.CalledProcessError:
        print("安装失败")
        return False
    
    # 等待服务启动
    time.sleep(5)
    return check_ollama_installed()

def print_models():
    print("可用模型列表：")
    for i, (name, size, rec) in enumerate(MODELS, 1):
        print(f"{i}. {name.ljust(15)} {size.ljust(8)} {rec}")

def get_model_choice():
    while True:
        try:
            choice = int(input("请选择要安装的模型编号 (1-7): "))
            if 1 <= choice <= 7:
                return choice - 1
            print("请输入1-7之间的数字")
        except ValueError:
            print("请输入有效的数字")

def is_model_installed(model_name):
    try:
        result = subprocess.run(
            ["ollama", "list"],
            capture_output=True,
            text=True,
            check=True
        )
        return model_name in result.stdout
    except subprocess.CalledProcessError:
        return False

def install_model(model_name, retries=2):
    for attempt in range(retries):
        try:
            print(f"正在安装模型 {model_name} (尝试 {attempt+1}/{retries})... 请耐心等待")
            subprocess.run(
                ["ollama", "run", model_name],
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True
            )
            return True
        except subprocess.CalledProcessError as e:
            print(f"安装失败: {e.output}")
    return False

def main():
    # 检查Ollama安装
    if not check_ollama_installed():
        print("Ollama未安装，开始安装...")
        while not install_ollama():
            print("正在重试安装...")
        print("Ollama安装成功!")
    
    # 打印模型列表
    print_models()
    
    # 获取用户选择
    choice = get_model_choice()
    model_name, size, rec = MODELS[choice]
    
    # 检查模型是否已安装
    if is_model_installed(model_name):
        print(f"模型 {model_name} 已经安装")
        return
    
    # 安装模型
    if install_model(model_name):
        print(f"模型 {model_name} 安装成功!")
    else:
        print(f"模型 {model_name} 安装失败，请检查网络或磁盘空间")

if __name__ == "__main__":
    main()